import pymysql
from random import choice
from SpiderProject.Spider.settings import MYSQL_CONFIG


class Saver(object):
    def __init__(self):
        self.db = pymysql.connect(**MYSQL_CONFIG)
        self.cursor = self.db.cursor(pymysql.cursors.DictCursor)
        self.__table = 'proxies'

    def add_data(self, proxy):
        select_sql = f'select * from {self.__table} where proxy=%s'
        self.cursor.execute(select_sql, (proxy,))
        result = self.cursor.fetchone()
        if not result:
            insert_sql = f'insert into {self.__table} (proxy, score) values (%s, %s)'
            try:
                self.cursor.execute(insert_sql, (proxy, 1))
                self.db.commit()
            except:
                self.db.rollback()
                print(proxy, '代理入库失败')

    def random_data(self):
        sql = f'select * from {self.__table} where score=1'
        self.cursor.execute(sql)
        all_proxy = self.cursor.fetchall()
        if all_proxy:
            result = choice(all_proxy)['proxy']
            update_sql = f'update {self.__table} set score=0 where proxy=%s'
            try:
                self.cursor.execute(update_sql, (result,))
                self.db.commit()
            except:
                self.db.rollback()
            return result
        else:
            print('当前没有可用代理')
            return None

    def delete_data(self):
        sql = f'delete from {self.__table} where score=0'
        try:
            self.cursor.execute(sql)
            self.db.commit()
        except:
            self.db.rollback()

    def count(self):
        sql = f'select * from {self.__table} where score=1'
        self.cursor.execute(sql)
        return self.cursor.rowcount


if __name__ == '__main__':
    a = Saver()
    sql = 'select * from proxies where score=1'
    a.cursor.execute(sql)
    print(a.cursor.fetchall())